package com.yungui.kafka.consumer.es;

import com.yungui.kafka.consumer.constant.EsConstant;
import com.yungui.kafka.consumer.entity.EsEntity;
import com.yungui.kafka.consumer.util.JsonUtil;
import org.apache.http.ConnectionClosedException;
import org.apache.http.conn.ConnectTimeoutException;
import org.elasticsearch.action.DocWriteResponse;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.rest.RestStatus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.IOException;
import java.net.ConnectException;
import java.net.SocketTimeoutException;
import java.util.*;
/**
 * @author laiaiqin
 * @since 2019/5/6
 */
@Service
public class EsCommonService {

    private final static Logger LOG = LoggerFactory.getLogger(EsCommonService.class);

    @Resource
    private RestHighLevelClient esClient;

    /**
     * 判断文档是否存在
     */
    public boolean isExistsDocumentById(String indexName, String id) throws IOException{

        GetRequest getRequest = new GetRequest(indexName, id);
        GetResponse getResponse = esClient.get(getRequest, RequestOptions.DEFAULT);
        boolean isExists = getResponse.isExists();
        LOG.info("判断{}索引文档{}是否存在，结果{}", indexName, id, isExists);
        return isExists;
    }

    /**
     * 创建文档
     */
    public boolean createDocumentByIdAndMap(String indexName, String id, Map<String,Object> sourceMap) throws IOException{

        IndexRequest indexRequest = new IndexRequest(indexName);
        indexRequest.id(id);
        indexRequest.source(sourceMap);
        IndexResponse indexResponse = esClient.index(indexRequest, RequestOptions.DEFAULT);
        boolean isSuccess = indexResponse.getResult() == DocWriteResponse.Result.CREATED;
        LOG.info("创建{}索引文档{}，结果{}", indexName, id, isSuccess);
        return isSuccess;
    }

    /**
     * 批量插入文档
     * @param index
     * @param list
     * @return
     */
    public BulkResponse insertBatch(String index, List<EsEntity> list) {
        BulkRequest request = new BulkRequest();
        for (EsEntity item : list) {
            String _json = JsonUtil.toJson(item.getData());
            String _id = item.getId();
            IndexRequest indexRequest = new IndexRequest(index).id(_id).source(_json, XContentType.JSON);
            request.add(indexRequest);
        }
        try {
            return esClient.bulk(request, RequestOptions.DEFAULT);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 批量插入文档
     * @param index
     * @param item
     * @return
     */
    public BulkResponse insert(String index, EsEntity item) {
        BulkRequest request = new BulkRequest();
        String _json = JsonUtil.toJson(item.getData());
        String _id = item.getId();
        IndexRequest indexRequest = new IndexRequest(index).id(_id).source(_json, XContentType.JSON);
        request.add(indexRequest);
        try {
            return esClient.bulk(request, RequestOptions.DEFAULT);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 更新文档
     */
    public boolean updateDocumentByIdAndMap(String indexName, String id, Map<String,Object> sourceMap) throws IOException{

        UpdateRequest updateRequest = new UpdateRequest(indexName, id);
        updateRequest.doc(sourceMap);
        UpdateResponse updateResponse = esClient.update(updateRequest,RequestOptions.DEFAULT);
        boolean isSuccess = updateResponse.getResult() == DocWriteResponse.Result.UPDATED;
        LOG.info("更新{}索引文档{}，结果{}", indexName, id, isSuccess);
        return isSuccess;
    }

    /**
     * 删除文档
     */
    public boolean deleteDocumentById(String indexName, String id) throws IOException{

        DeleteRequest deleteRequest = new DeleteRequest(indexName,id);
        DeleteResponse deleteResponse = esClient.delete(deleteRequest,RequestOptions.DEFAULT);
        boolean isSuccess = deleteResponse.getResult() == DocWriteResponse.Result.DELETED;
        LOG.info("删除{}索引文档{}，结果{}", indexName, id, isSuccess);
        return isSuccess;
    }


    /**
     * ES批量操作
     *
     * @param failureList
     * @param bulkRequest
     * @return
     * @throws Exception
     */
    private void processBulk(List<BulkItemResponse.Failure> failureList,
                             BulkRequest bulkRequest, Boolean refresh, int retryTimes) throws Exception {
        try {
            if (refresh) {
                //当refresh=wait_for请求返回时，它的更改对于搜索是可见的
                bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
            }
            BulkResponse bulkResponse = esClient.bulk(bulkRequest, RequestOptions.DEFAULT);
            Boolean neeRetry = false;
            if (bulkResponse.hasFailures()) {
                for (BulkItemResponse bulkItemResponse : bulkResponse) {
                	if(bulkItemResponse.isFailed()) {
                		if(RestStatus.TOO_MANY_REQUESTS.equals(bulkItemResponse.getFailure().getStatus())) {
                			neeRetry = true;
                			failureList.add(bulkItemResponse.getFailure());
                		} else {
                            LOG.error("index name is " + bulkItemResponse.getIndex());
                            LOG.error(bulkItemResponse.getFailureMessage());
                		}
                	}
                }
            }
            if (neeRetry && retryTimes < EsConstant.REJECT_MAX_RETRY_TIMES) {
                //EsRejectedExecutionException 线程池处理不过来，等待15秒，重试
                Thread.sleep(EsConstant.REJECT_SLEEP);
                LOG.error("catch EsRejectedExecutionException, retry bulk index request.");
                failureList.clear();
                processBulk(failureList, bulkRequest, refresh, retryTimes + 1);
            }
        } catch (ConnectException ce) {
            processConnectFail(ce);
        } catch (ConnectTimeoutException ctc) {
            processConnectFail(ctc);
        } catch (SocketTimeoutException ste) {
            processConnectFail(ste);
        } catch (ConnectionClosedException cce) {
            processConnectFail(cce);
        }
    }

    /**
     * 出现ES连接异常时的处理
     *
     * @param e
     */
    public void processConnectFail(Exception e) throws Exception {
        throw e;
    }
}
